import numpy as np
from astropy.time import Time
from astropy.coordinates import Longitude, Latitude, Angle
from astropy import coordinates
import astropy.units as u
import utils



class InitialOrbit:
    def __init__(self, data_list, time_utc):
        self.data_list = data_list
        self.time_utc = time_utc
        self.lambda_mu_nu_list = list()
        self.F = 0
        self.F_z = 0
        self.G = 0
        self.G_z = 0
        self.tau = 0
        self.R = np.zeros(3)
        self.r = np.zeros(3)
        self.r_dot = np.zeros(3)
        self.t_0 = 0

    def get_site_coord(self):

        site_lon = 120 * u.deg
        site_lat = 36 * u.deg
        site_height = 40 * u.m
        # 测站是ITRS坐标系
        site_coord_ITRS = coordinates.EarthLocation.from_geodetic(site_lon, site_lat,
                                                                  site_height).get_itrs().cartesian.xyz
        transform_matrix = utils.CoordinatesTransformEquinoxBase(self.time_utc)
        # 再转到GCRS坐标系，这个就是R
        site_coord_GCRS = (transform_matrix.precession_matrix() @ transform_matrix.nutation_matrix() @ transform_matrix.earth_rotation_matrix()).T @ site_coord_ITRS
        self.R = site_coord_GCRS.value / a_e
        self.r = self.R * 2

    def get_lambda_mu_nu(self):
        self.t_0 = data_list[0][0].to_value('s') / time_unit
        for data in self.data_list:
            data_day_second, azimuth, altitude = data
            data_time = self.time_utc + data_day_second
            transform_matrix = utils.CoordinatesTransformEquinoxBase(data_time)
            matrix_GR = transform_matrix.nutation_matrix() @ transform_matrix.precession_matrix()
            sg = data_time.sidereal_time('apparent', 'greenwich').to_value('hourangle') * np.pi / 12
            matrix_ZR = utils.rotation_matrix('z', np.pi - sg) @ utils.rotation_matrix('y', np.pi / 2 - np.pi / 5)
            lambda_mu_nu = (matrix_GR.T @
                            matrix_ZR @
                            np.array([np.cos(altitude) * np.cos(azimuth),
                                      -np.cos(altitude) * np.sin(azimuth),
                                      np.sin(altitude)]))
            self.lambda_mu_nu_list.append((data_day_second.to_value('s') / time_unit - self.t_0, lambda_mu_nu))
        self.F = 1
        self.F_z = 1
        self.tau = data_list[0][0] / time_unit
        self.tau = self.tau.to_value('s')
        self.G = self.tau
        self.G_z = self.tau

    def update_solution(self):
        A = np.zeros((3 * len(self.lambda_mu_nu_list), 6))
        B = np.zeros(3 * len(self.lambda_mu_nu_list))
        X, Y, Z = self.R
        for i in range(len(self.lambda_mu_nu_list)):
            delta_t, lmda_mu_nu = self.lambda_mu_nu_list[i]
            if i != 0:
                self.update_F_G(delta_t)
            lmda, mu, nu = lmda_mu_nu
            A[3 * i + 0] = np.array([self.F * nu, 0, -self.F_z * lmda, self.G * nu, 0, -self.G_z * lmda])
            A[3 * i + 1] = np.array([0, self.F * nu, -self.F_z * mu, 0, self.G * nu, -self.G_z * mu])
            A[3 * i + 2] = np.array([self.F * mu, -self.F * lmda, 0, self.G * mu, -self.G * lmda, 0])
            B[3 * i + 0] = nu * X - lmda * Z
            B[3 * i + 1] = nu * Y - mu * Z
            B[3 * i + 2] = mu * X - lmda * Y

        solution = np.linalg.lstsq(A, B, rcond=None)
        self.r = solution[0][:3]
        self.r_dot = solution[0][3:]
    def update_F_G(self, this_tau):
        r = np.linalg.norm(self.r)
        z = self.r[2]
        z_dot = self.r_dot[2]
        u = np.zeros(10)
        tau = np.zeros(8)
        u[0] = 1
        tau[0] = 1
        for i in range(1, 10):
            u[i] = 1 / (r ** i)
        for i in range(1, 8):
            tau[i] = this_tau ** i
        sigma = np.dot(self.r, self.r_dot)
        v_0_2 = np.linalg.norm(self.r_dot) ** 2
        F = ((1 + tau[2] / 2 * (-1 * u[3])
                + tau[3] * 1/6 * (3 * u[5] * sigma + 10 * u[7] * z * z_dot))
                + tau[4] / 24 * (u[5] * (3 * v_0_2 - 2 * u[1] - 15 * u[2] * (sigma ** 2))
                                 - 5 * u[7] * (7 * u[2] * (z ** 2) - 1) * v_0_2
                                 + 10 * u[7] * (z_dot ** 2)
                                 + 35 * u[9] * (9 * u[2] * (z ** 2) - 1) * (sigma ** 2)
                                 - 140 * u[9] * sigma * z * z_dot
                                )
                + tau[5] / 120 * u[7] * (15 * sigma * (-3 * v_0_2 + 2 * u[1] + 7 * u[2] * (sigma ** 2)))
                + tau[6] / 720 * u[7] * (u[2] * (sigma ** 2) * (630 * v_0_2 - 420 * u[1] - 945 * u[2] * (sigma ** 2))
                                         - (22 * u[2] - 66 * u[1] * v_0_2 + 45 * (v_0_2 ** 2))))

        F_z = F

        G = (tau[1] + tau[3] / 6 * (-u[3]) + tau[4] / 24 * (6 * u[5] * sigma) +
             tau[5] / 120 * u[5] * (9 * v_0_2 - 8 * u[1] - 45 * u[2] * (sigma ** 2)) +
             tau[6] / 720 * u[7] * (30 * sigma * (-6 * v_0_2 + 5 * u[1] + 14 * u[2] * (sigma ** 2))))

        G_z = G
        self.F = F
        self.F_z = F_z
        self.G = G
        self.G_z = G_z


a_e = 6378137
GE = 3986004418e5
time_unit = np.sqrt(a_e ** 3 / GE)

time_str = '2002-03-30T00:00:00'
utc_time = Time(time_str, scale='utc') + 28470 * u.s
data_list = utils.read_txt('data.txt')

initial_orbit = InitialOrbit(data_list, utc_time)
initial_orbit.get_site_coord()
initial_orbit.get_lambda_mu_nu()
# initial_orbit.update_F_G(initial_orbit.tau)
# 最开始设置了100次循环，发现差不多第50次左右这个F和G就变化的特别小了。
for i in range(500):
    initial_orbit.update_solution()
    # initial_orbit.update_F_G()

orbit_elements = utils.get_orbit_elements(initial_orbit.r * a_e, initial_orbit.r_dot * a_e / time_unit)
print(orbit_elements)